{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Building a Search Engine with BERT and TensorFlow",
      "version": "0.3.2",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wy8WsIgu1qE7",
        "colab_type": "text"
      },
      "source": [
        "# Building a Search Engine with BERT and TensorFlow"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NR9bffZha4Da",
        "colab_type": "text"
      },
      "source": [
        "In this experiment, we will use a pre-trained BERT model checkpoint to build a general-purpose text feature extractor.\n",
        "\n",
        "These things are sometimes referred to as  [Natural Language Understanding](https://en.wikipedia.org/wiki/Natural-language_understanding) (NLU) modules, because the features they extract are relevant for a wide array of downstream NLP tasks.\n",
        "\n",
        "One use for these features is in [instance-based learning](https://en.wikipedia.org/wiki/Instance-based_learning), which relies on computing the similarity of the query to the training samples.\n",
        "\n",
        "We will illustrate this by building a simple Information Retrieval system using the BERT NLU module for feature extraction.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DepWa1CZEv1O",
        "colab_type": "text"
      },
      "source": [
        "The plan for this experiment is:\n",
        "1. getting the pre-trained BERT model checkpoint\n",
        "2. extracting a sub-graph optimized for inference\n",
        "3. creating a feature extractor with tf.Estimator\n",
        "4. implementing an Information Retrieval engine\n",
        "5. accelerating search queries with math\n",
        "6. building a movie recommendation system"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gRwsCW6w7Ztg",
        "colab_type": "text"
      },
      "source": [
        "## Step 1: getting the pre-trained model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "frHPPNtB7XQJ",
        "colab_type": "text"
      },
      "source": [
        "We start with a pre-trained english BERT-base model checkpoint. \n",
        "\n",
        "For configuring and optimizing the graph for inference we will use [bert-as-a-service](https://github.com/hanxiao/bert-as-service) repository, which allows for serving BERT models for remote clients over TCP.\n",
        "\n",
        "Having a single BERT-server is definitely beneficial in many environments. However,  in this part of the experiment we will focus on creating a local (in-process) feature extractor. This is useful if one wishes to avoid additional latency and potential failure modes introduced by a client-server architecture.\n",
        "\n",
        "Now, let us download the model and install the package."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WYYMcLgNa2cy",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!wget https://storage.googleapis.com/bert_models/2018_10_18/uncased_L-12_H-768_A-12.zip\n",
        "!unzip uncased_L-12_H-768_A-12.zip\n",
        "!pip install bert-serving-server --no-deps"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AvvFoXCR1zeK",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SZQHFXL4Afgg",
        "colab_type": "text"
      },
      "source": [
        "## Step 2: optimizing the inference graph\n",
        "Normally, to modify the model graph we would have to do some low-level TensorFlow programming. \n",
        "\n",
        "However, thanks to bert-as-a-service, we can configure the inference graph using a simple CLI interface."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "m0NnkrTs7GVe",
        "colab_type": "text"
      },
      "source": [
        "There are a couple of parameters in the below snippet too look out for.\n",
        "\n",
        "For each text sample, BERT-base model encoding layers output a tensor of shape **[sequence_len, encoder_dim],** with one vector per input token. To obtain a fixed representation, we need to apply some sort of pooling.\n",
        "\n",
        "**POOL_STRAT** parameter defines the pooling strategy applied to the  **POOL_LAYER** encoding layer. The default value **REDUCE_MEAN** averages the vectors for all tokens in a sequence. This strategy works best for most sentence-level tasks, when the model is not fine-tuned. Another option is NONE, in which case no pooling is applied at all. This is useful for word-level tasks such as Named Entity Recognition or POS tagging. For a detailed discussion of other options check out the Han Xiao's [blog post.](https://hanxiao.github.io/2019/01/02/Serving-Google-BERT-in-Production-using-Tensorflow-and-ZeroMQ/)\n",
        "\n",
        "**SEQ_LEN** affects the maximum length of sequences processed by the model. Smaller values increase the model inference speed almost linearly."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FLJUBvG72_wo",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import os\n",
        "import tensorflow as tf\n",
        "sesh = tf.InteractiveSession()\n",
        "\n",
        "from bert_serving.server.graph import optimize_graph\n",
        "from bert_serving.server.helper import get_args_parser\n",
        "\n",
        "# input dir\n",
        "MODEL_DIR = '/content/uncased_L-12_H-768_A-12' #@param {type:\"string\"}\n",
        "# output dir\n",
        "GRAPH_DIR = '/content/graph/' #@param {type:\"string\"}\n",
        "# output filename\n",
        "GRAPH_OUT = 'extractor.pbtxt' #@param {type:\"string\"}\n",
        "\n",
        "POOL_STRAT = 'REDUCE_MEAN' #@param ['REDUCE_MEAN', 'REDUCE_MAX', \"NONE\"]\n",
        "POOL_LAYER = '-2' #@param {type:\"string\"}\n",
        "SEQ_LEN = '256' #@param {type:\"string\"}\n",
        "\n",
        "tf.gfile.MkDir(GRAPH_DIR)\n",
        "\n",
        "parser = get_args_parser()\n",
        "carg = parser.parse_args(args=['-model_dir', MODEL_DIR,\n",
        "                               '-graph_tmp_dir', GRAPH_DIR,\n",
        "                               '-max_seq_len', str(SEQ_LEN),\n",
        "                               '-pooling_layer', str(POOL_LAYER),\n",
        "                               '-pooling_strategy', POOL_STRAT])\n",
        "\n",
        "tmp_name, config = optimize_graph(carg)\n",
        "graph_fout = os.path.join(GRAPH_DIR, GRAPH_OUT)\n",
        "\n",
        "tf.gfile.Rename(\n",
        "    tmp_name,\n",
        "    graph_fout,\n",
        "    overwrite=True\n",
        ")\n",
        "print(\"\\nSerialized graph to {}\".format(graph_fout))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "t0iyosxG8zLo",
        "colab_type": "text"
      },
      "source": [
        "Running the above snippet will put the BERT model graph and weights from  **MODEL_DIR** into a GraphDef object which will be serialized to a pbtxt file at **GRAPH_OUT**. The file will be smaller than the original model because the nodes and variables required for training will be removed. This results in a quite portable solution: for example the english base model only takes 389 MB after exporting."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3zqy0cdx7Ty0",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "A8zuhd7k9v7W",
        "colab_type": "text"
      },
      "source": [
        "## Step 3: creating a feature extractor"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ICfRIpI4-K8t",
        "colab_type": "text"
      },
      "source": [
        "Now, we will use the serialized graph to build a feature extractor using the tf.Estimator API. We will need to define two things: **input_fn** and **model_fn**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RmgCepvNsJv5",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import logging\n",
        "import numpy as np\n",
        "\n",
        "from tensorflow.python.estimator.estimator import Estimator\n",
        "from tensorflow.python.estimator.run_config import RunConfig\n",
        "from tensorflow.python.estimator.model_fn import EstimatorSpec\n",
        "from tensorflow.keras.utils import Progbar\n",
        "\n",
        "from bert_serving.server.bert.tokenization import FullTokenizer\n",
        "from bert_serving.server.bert.extract_features import convert_lst_to_features\n",
        "\n",
        "\n",
        "log = logging.getLogger('tensorflow')\n",
        "log.setLevel(logging.INFO)\n",
        "log.handlers = []"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qIEZPsu-kOzx",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "GRAPH_PATH = \"/content/graph/extractor.pbtxt\" #@param {type:\"string\"}\n",
        "VOCAB_PATH = \"/content/uncased_L-12_H-768_A-12/vocab.txt\" #@param {type:\"string\"}\n",
        "\n",
        "SEQ_LEN = 256 #@param {type:\"integer\"}"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FhsEwpOF8LNs",
        "colab_type": "text"
      },
      "source": [
        "## input_fn"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PNJ2M3W8-kWw",
        "colab_type": "text"
      },
      "source": [
        "**input_fn** manages getting the data into the model. That includes executing the whole text preprocessing pipeline and preparing a feed_dict for BERT. \n",
        "\n",
        "First, each text sample is converted into a tf.Example instance containing the necessary features listed in **INPUT_NAMES**. The bert_tokenizer object contains  the WordPiece vocabulary and performs the text preprocessing. After that the examples are re-grouped by feature name in a **feed_dict**."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nULWk9-fkO5I",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "INPUT_NAMES = ['input_ids', 'input_mask', 'input_type_ids']\n",
        "bert_tokenizer = FullTokenizer(VOCAB_PATH)\n",
        "\n",
        "def build_feed_dict(texts):\n",
        "    \n",
        "    text_features = list(convert_lst_to_features(\n",
        "        texts, SEQ_LEN, SEQ_LEN, \n",
        "        bert_tokenizer, log, False, False))\n",
        "\n",
        "    target_shape = (len(texts), -1)\n",
        "\n",
        "    feed_dict = {}\n",
        "    for iname in INPUT_NAMES:\n",
        "        features_i = np.array([getattr(f, iname) for f in text_features])\n",
        "        features_i = features_i.reshape(target_shape).astype(\"int32\")\n",
        "        feed_dict[iname] = features_i\n",
        "\n",
        "    return feed_dict"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oVDwr_cR_V0U",
        "colab_type": "text"
      },
      "source": [
        "tf.Estimators have a fun feature which makes them re-build and re-initialize the whole computational graph at each call to the predict function. \n",
        "\n",
        "So, in order to avoid the overhead, to the predict function we will pass a generator, which will yield the features to the model in a never-ending loop. Ha-ha."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ljud_8WXIMeC",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def build_input_fn(container):\n",
        "    \n",
        "    def gen():\n",
        "        while True:\n",
        "          try:\n",
        "            yield build_feed_dict(container.get())\n",
        "          except:\n",
        "            yield build_feed_dict(container.get())\n",
        "\n",
        "    def input_fn():\n",
        "        return tf.data.Dataset.from_generator(\n",
        "            gen,\n",
        "            output_types={iname: tf.int32 for iname in INPUT_NAMES},\n",
        "            output_shapes={iname: (None, None) for iname in INPUT_NAMES})\n",
        "    return input_fn\n",
        "\n",
        "class DataContainer:\n",
        "  def __init__(self):\n",
        "    self._texts = None\n",
        "  \n",
        "  def set(self, texts):\n",
        "    if type(texts) is str:\n",
        "      texts = [texts]\n",
        "    self._texts = texts\n",
        "    \n",
        "  def get(self):\n",
        "    return self._texts"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UGgF6U9B81Zf",
        "colab_type": "text"
      },
      "source": [
        "## model_fn"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9Ck6x21UBIOA",
        "colab_type": "text"
      },
      "source": [
        "model_fn contains the specification of the model. In our case, it is loaded from the pbtxt file we saved in the previous step. \n",
        "\n",
        "The features are mapped explicitly to the corresponding input nodes with input_map."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Mr4gQqBKS1rp",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def model_fn(features, mode):\n",
        "    with tf.gfile.GFile(GRAPH_PATH, 'rb') as f:\n",
        "        graph_def = tf.GraphDef()\n",
        "        graph_def.ParseFromString(f.read())\n",
        "        \n",
        "    output = tf.import_graph_def(graph_def,\n",
        "                                 input_map={k + ':0': features[k] for k in INPUT_NAMES},\n",
        "                                 return_elements=['final_encodes:0'])\n",
        "\n",
        "    return EstimatorSpec(mode=mode, predictions={'output': output[0]})\n",
        "  \n",
        "estimator = Estimator(model_fn=model_fn)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4l4EPmwtBadU",
        "colab_type": "text"
      },
      "source": [
        "Now we have everything we need to perform inference. Let's do this!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "37oYNYf6MX1R",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def batch(iterable, n=1):\n",
        "    l = len(iterable)\n",
        "    for ndx in range(0, l, n):\n",
        "        yield iterable[ndx:min(ndx + n, l)]\n",
        "\n",
        "def build_vectorizer(_estimator, _input_fn_builder, batch_size=128):\n",
        "  container = DataContainer()\n",
        "  predict_fn = _estimator.predict(_input_fn_builder(container), yield_single_examples=False)\n",
        "  \n",
        "  def vectorize(text, verbose=False):\n",
        "    x = []\n",
        "    bar = Progbar(len(text))\n",
        "    for text_batch in batch(text, batch_size):\n",
        "      container.set(text_batch)\n",
        "      x.append(next(predict_fn)['output'])\n",
        "      if verbose:\n",
        "        bar.add(len(text_batch))\n",
        "      \n",
        "    r = np.vstack(x)\n",
        "    return r\n",
        "  \n",
        "  return vectorize"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SqTg_labMb3E",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "bert_vectorizer = build_vectorizer(estimator, build_input_fn)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "etpTDXf-MX5H",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "bert_vectorizer(64*['sample text']).shape"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SaxAbqlRIXCZ",
        "colab_type": "text"
      },
      "source": [
        "Consecutive calls will not rebuild the model."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NLs9HMj5IVbU",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "bert_vectorizer(256*['sample text']).shape"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Cl6PSKjFZkfS",
        "colab_type": "text"
      },
      "source": [
        "## Step 4: exploring vector space with Projector"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nlQnReadQ6jQ",
        "colab_type": "text"
      },
      "source": [
        "A standalone version of BERT feature extractor is available in the [repository](https://github.com/gaphex/bert_experimental)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BNjUPnOaF43_",
        "colab_type": "text"
      },
      "source": [
        "Now it's time for a demonstration!\n",
        "\n",
        "Using the vectorizer we will generate embeddings for articles from the Reuters-21578 benchmark corpus.\n",
        "To visualise and explore the embedding vector space in 3D we will use a dimensionality reduction technique called [T-SNE](https://distill.pub/2016/misread-tsne/).\n",
        "\n",
        "Lets get the article embeddings first."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VoYOSce16iBG",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import nltk\n",
        "from nltk.corpus import reuters\n",
        "\n",
        "nltk.download(\"reuters\")\n",
        "nltk.download(\"punkt\")\n",
        "\n",
        "max_samples = 256\n",
        "categories = ['wheat', 'tea', 'strategic-metal', \n",
        "              'housing', 'money-supply', 'fuel']\n",
        "\n",
        "S, X, Y = [], [], []\n",
        "\n",
        "for category in categories:\n",
        "  print(category)\n",
        "  \n",
        "  sents = reuters.sents(categories=category)\n",
        "  sents = [' '.join(sent) for sent in sents][:max_samples]\n",
        "  X.append(bert_vectorizer(sents, verbose=True))\n",
        "  Y += [category] * len(sents)\n",
        "  S += sents\n",
        "  \n",
        "X = np.vstack(X) \n",
        "X.shape"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PIIHDuuLLydo",
        "colab_type": "text"
      },
      "source": [
        "The interactive visualization of generated embeddings is available on the [Embedding Projector](https://projector.tensorflow.org/?config=https://gist.githubusercontent.com/gaphex/7262af1e151957b1e7c638f4922dfe57/raw/3b946229fc58cbefbca2a642502cf51d4f8e81c5/reuters_proj_config.json). \n",
        "\n",
        "From the link you can run T-SNE yourself, or load a checkpoint using the bookmark in lower-right corner (loading works only on Chrome).\n",
        "\n",
        "To reproduce the input files used for this visualization, run the code below. Then, download the files to your machine and upload to Projector\n",
        "\n",
        "(you can dowload files from the menu opened by the \">\" button in the upper-left)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "d4taY-zv6iMt",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "with open(\"embeddings.tsv\", \"w\") as fo:\n",
        "  for x in X.astype('float16'):\n",
        "    line = \"\\t\".join([str(v) for v in x])\n",
        "    fo.write(line + \"\\n\")\n",
        "\n",
        "with open(\"metadata.tsv\", \"w\") as fo:\n",
        "  fo.write(\"Label\\tSentence\\n\")\n",
        "  for y, s in zip(Y, S):\n",
        "    fo.write(\"{}\\t{}\\n\".format(y, s))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9TwkuXXDGHGI",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ayA6eJTpJCN8",
        "colab_type": "text"
      },
      "source": [
        "Here's what I captured using the Projector."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lEXazXzuGHRI",
        "colab_type": "code",
        "outputId": "c0bc5780-7143-4771-9085-a54f00cb186a",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 653
        }
      },
      "source": [
        "from IPython.display import HTML\n",
        "\n",
        "HTML(\"\"\"\n",
        "<video width=\"900\" height=\"632\" controls>\n",
        "  <source src=\"https://storage.googleapis.com/bert_resourses/reuters_tsne_hd.mp4\" type=\"video/mp4\">\n",
        "</video>\n",
        "\"\"\")"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/html": [
              "\n",
              "<video width=\"900\" height=\"632\" controls>\n",
              "  <source src=\"https://storage.googleapis.com/bert_resourses/reuters_tsne_hd.mp4\" type=\"video/mp4\">\n",
              "</video>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 22
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "q97jPimBuOxY",
        "colab_type": "text"
      },
      "source": [
        "## Step 5: building a nearest neighbour search engine"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Gfibka1lua4Z",
        "colab_type": "text"
      },
      "source": [
        "Now, let's say we have a knowledge base of 50k test samples, and we need to answer queries based on this data, fast. \n",
        "\n",
        "How do we retrieve the sample, most similar to a query, from a text database? The answer is nearest neighbour search.\n",
        "\n",
        "Formally, the search problem we will be solving is defined as follows:\n",
        "\n",
        "*given a set of points **S** in vector space **M**, and a query point **Q** ∈ **M**, find the closest point in **S** to **Q**.* \n",
        "\n",
        "There are multiple ways to define 'closest' in vector space, but we will use Euclidean distance.\n",
        "\n",
        "So, to build an Information Retrieval system, we will follow these steps:\n",
        "\n",
        "1.   Vectorize all samples from the knowledge base - that gives **S**\n",
        "2.   Vectorize the query - that gives **Q**\n",
        "3.   Compute euclidean distances **D** between **Q** and **S**\n",
        "4.   Sort **D** in ascending order , providing indices of the most similar samples\n",
        "5.   Retrieve labels for said samples from the knowledge base"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Eaa6KoHLu0fm",
        "colab_type": "text"
      },
      "source": [
        "To make the simple matter of implementing this a bit more exciting, we will do it in pure TensorFlow."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XUnV8LDoIhbL",
        "colab_type": "text"
      },
      "source": [
        "First we define the placeholders for Q and S"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SA6JCVwa8IhT",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "graph = tf.Graph()\n",
        "sess = tf.InteractiveSession(graph=graph)\n",
        "\n",
        "dim = X.shape[1]\n",
        "\n",
        "Q = tf.placeholder(\"float\", [dim])\n",
        "S = tf.placeholder(\"float\", [None, dim])"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZIRo06u0Iudw",
        "colab_type": "text"
      },
      "source": [
        "Define distance computation"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WmLDDjbhIu-1",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "squared_distance = tf.reduce_sum(tf.pow(Q - S, 2), reduction_indices=1)\n",
        "distance = tf.sqrt(squared_distance)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cl2Laz1dI5Qv",
        "colab_type": "text"
      },
      "source": [
        "Finally, get the most similar indices"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZEHys2u1Uqsc",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "top_k = 10\n",
        "\n",
        "top_neg_dists, top_indices = tf.math.top_k(tf.negative(distance), k=top_k)\n",
        "top_dists = tf.negative(top_neg_dists)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oGuFRZUvFnvV",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NUW78Ns-JH4I",
        "colab_type": "text"
      },
      "source": [
        "Sanity check"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "C7doXw3Z-kdB",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from sklearn.metrics.pairwise import euclidean_distances"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TOPb3SRzFqqT",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "top_indices.eval({Q:X[0], S:X})"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tGrb-FahFqwV",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "np.argsort(euclidean_distances(X[:1], X)[0])[:10]"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eStMaiQfrFFb",
        "colab_type": "text"
      },
      "source": [
        "## Step 6: accelerating search"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hXJIG0RnrFIo",
        "colab_type": "text"
      },
      "source": [
        "Now that we have a basic retrieval engine set up, the question is: can we make it run faster? \n",
        "With a tiny bit of math, we can.\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IahEjvops6I4",
        "colab_type": "text"
      },
      "source": [
        "For a pair of vectors p and q the euclidean distance is defined as follows:\n",
        "\n",
        "$$ d(\\mathbf{p}, \\mathbf{q}) = d(\\mathbf{q}, \\mathbf{p}) = \\sqrt{(p_{1}-q_{1})^{2}+(p_{2}-q_{2})^{2}+ ... + (p_{n}-q_{n})^{2}} = \\sqrt{\\sum_{i=1}^{n}(p_{i}-q_{i})^{2}}$$\n",
        "\n",
        "Which is exactly how we did compute it in step 4."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NBiXZZJTuWic",
        "colab_type": "text"
      },
      "source": [
        "However, since p and q are vectors, we can expand and rewrite this:\n",
        "\n",
        "$$ d(\\mathbf{p}, \\mathbf{q})^{2} = \\langle \\mathbf{p} - \\mathbf{q}, \\mathbf{p} - \\mathbf{q} \\rangle  =\\langle \\mathbf{p} \\mathbf{p} \\rangle - 2 \\langle \\mathbf{p} \\mathbf{q} \\rangle + \\langle \\mathbf{q} \\mathbf{q} \\rangle $$\n",
        "\n",
        "where ⟨…⟩ denotes inner product."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3lwHGnkIyZUD",
        "colab_type": "text"
      },
      "source": [
        "In TensorFlow this can be written as follows:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4yDVIZoQCGJY",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "Q = tf.placeholder(\"float\", [dim])\n",
        "S = tf.placeholder(\"float\", [None, dim])\n",
        "\n",
        "Qr = tf.reshape(Q, (1, -1))\n",
        "\n",
        "PP = tf.keras.backend.batch_dot(S, S, axes=1)\n",
        "QQ = tf.matmul(Qr, tf.transpose(Qr))\n",
        "PQ = tf.matmul(S, tf.transpose(Qr))\n",
        "\n",
        "distance = PP - 2 * PQ + QQ\n",
        "distance = tf.sqrt(tf.reshape(distance, (-1,)))\n",
        "\n",
        "top_neg_dists, top_indices = tf.math.top_k(tf.negative(distance), k=top_k)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OEWw7LOLZdCc",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "top_indices.eval({Q:X[0], S:X})"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "M49F0Le6TuH3",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PFFe18cL1EC8",
        "colab_type": "text"
      },
      "source": [
        "By the way, in the formula above PP and QQ are actually squared [L2 norms](https://en.wikipedia.org/wiki/Norm_(mathematics)#Euclidean_norm) of the respective vectors. If both vectors are L2 normalized, then PP = QQ = 1. \n",
        "\n",
        "That gives an interesting relation between inner product and euclidean distance for L2 normalized data:\n",
        "\n",
        "$$d(\\mathbf{p}, \\mathbf{q})^{2} = 2 * (1 - \\langle \\mathbf{p} \\mathbf{q} \\rangle )$$"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gu9KmHws1l_f",
        "colab_type": "text"
      },
      "source": [
        "However, doing L2 normalization discards the information about vector magnetude, which is undesirable in many cases. \n",
        "\n",
        "Instead, we may notice that as long as the knowledge base does not change, PP, it's vector norm, also remains constant. So, instead of recomputing it every time, we can just do it once and then use the precomputed result, further accelerating the distance computation."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ESN3ZFxw346p",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "Q = tf.placeholder(\"float\", [dim])\n",
        "S = tf.placeholder(\"float\", [None, dim])\n",
        "S_norm = tf.placeholder(\"float\", [None, 1])\n",
        "\n",
        "Qr = tf.reshape(Q, (1, -1))\n",
        "\n",
        "PP = S_norm\n",
        "QQ = tf.matmul(Qr, tf.transpose(Qr))\n",
        "PQ = tf.matmul(S, tf.transpose(Qr))\n",
        "\n",
        "distance = PP - 2 * PQ + QQ\n",
        "distance = tf.sqrt(tf.reshape(distance, (-1,)))\n",
        "\n",
        "top_neg_dists, top_indices = tf.math.top_k(tf.negative(distance), k=top_k)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pXEeJHgB4JNw",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "X_norm = np.sum(X**2, axis=1).reshape((-1,1))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MPn1RWt58kXk",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "top_indices.eval({Q:X[0], S:X, S_norm: X_norm})"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Z1P_l80-TN_z",
        "colab_type": "text"
      },
      "source": [
        "Now, let us put it all together:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "j1U4xy5sETxf",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class L2Retriever:\n",
        "    def __init__(self, dim, top_k=3, use_norm=False, use_gpu=True):\n",
        "        self.dim = dim\n",
        "        self.top_k = top_k\n",
        "        self.use_norm = use_norm\n",
        "        config = tf.ConfigProto(\n",
        "            device_count={'GPU': (1 if use_gpu else 0)}\n",
        "        )\n",
        "        self.session = tf.Session(config=config)\n",
        "        \n",
        "        self.norm = None\n",
        "        self.query = tf.placeholder(\"float\", [self.dim])\n",
        "        self.kbase = tf.placeholder(\"float\", [None, self.dim])\n",
        "        \n",
        "        self.build_graph()\n",
        "\n",
        "    def build_graph(self):\n",
        "      \n",
        "        if self.use_norm:\n",
        "            self.norm = tf.placeholder(\"float\", [None, 1])\n",
        "\n",
        "        distance = dot_l2_distances(self.kbase, self.query, self.norm)\n",
        "        top_neg_dists, top_indices = tf.math.top_k(tf.negative(distance), k=self.top_k)\n",
        "        top_dists = tf.negative(top_neg_dists)\n",
        "\n",
        "        self.top_distances = top_dists\n",
        "        self.top_indices = top_indices\n",
        "\n",
        "    def predict(self, kbase, query, norm=None):\n",
        "        query = np.squeeze(query)\n",
        "        feed_dict = {self.query: query, self.kbase: kbase}\n",
        "        if self.use_norm:\n",
        "          feed_dict[self.norm] = norm\n",
        "        \n",
        "        I, D = self.session.run([self.top_indices, self.top_distances],\n",
        "                                feed_dict=feed_dict)\n",
        "        \n",
        "        return I, D\n",
        "      \n",
        "def dot_l2_distances(kbase, query, norm=None):\n",
        "    query = tf.reshape(query, (1, -1))\n",
        "    \n",
        "    if norm is None:\n",
        "      XX = tf.keras.backend.batch_dot(kbase, kbase, axes=1)\n",
        "    else:\n",
        "      XX = norm\n",
        "    YY = tf.matmul(query, tf.transpose(query))\n",
        "    XY = tf.matmul(kbase, tf.transpose(query))\n",
        "    \n",
        "    distance = XX - 2 * XY + YY\n",
        "    distance = tf.sqrt(tf.reshape(distance, (-1,)))\n",
        "    \n",
        "    return distance"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aoMf-lWrx_jt",
        "colab_type": "text"
      },
      "source": [
        "## Example: movie recommendation system"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TJbkrOeu62lz",
        "colab_type": "text"
      },
      "source": [
        "For this example we will use a dataset of movie summaries from IMDB. \n",
        "\n",
        "We will vectorize the reviews and use the search engine to find nearest neighbours in movie space! \n",
        "\n",
        "That will allow us to build a simple movie recommendation system, which will suggest similar movies based on plot features.\n",
        "\n",
        "First let's download and prepare the dataset."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GKUjZt61EUiV",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import pandas as pd\n",
        "import json\n",
        "\n",
        "# download and build the dataset \n",
        "!wget http://www.cs.cmu.edu/~ark/personas/data/MovieSummaries.tar.gz\n",
        "!tar -xvzf MovieSummaries.tar.gz\n",
        "\n",
        "plots_df = pd.read_csv('MovieSummaries/plot_summaries.txt', sep='\\t', header=None)\n",
        "meta_df = pd.read_csv('MovieSummaries/movie.metadata.tsv', sep='\\t', header=None)\n",
        "\n",
        "plot = {}\n",
        "metadata = {}\n",
        "movie_data = {}\n",
        "\n",
        "for movie_id, movie_plot in plots_df.values:\n",
        "  plot[movie_id] = movie_plot \n",
        "\n",
        "for movie_id, movie_name, movie_genre in meta_df[[0,2,8]].values:\n",
        "  genre = list(json.loads(movie_genre).values())\n",
        "  if len(genre):\n",
        "    metadata[movie_id] = {\"name\": movie_name,\n",
        "                          \"genre\": genre}\n",
        "    \n",
        "for movie_id in set(plot.keys())&set(metadata.keys()):\n",
        "  movie_data[metadata[movie_id]['name']] = {\"genre\": metadata[movie_id]['genre'],\n",
        "                                            \"plot\": plot[movie_id]}\n",
        "\n",
        "X, Y, names = [], [], []\n",
        "\n",
        "for movie_name, movie_meta in movie_data.items():\n",
        "  X.append(movie_meta['plot'])\n",
        "  Y.append(movie_meta['genre'])\n",
        "  names.append(movie_name)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "z4iFv1ESUX5D",
        "colab_type": "text"
      },
      "source": [
        "Vectorize movie plots with BERT NLU. This will take a while."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "smgRnX2-Gviw",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "X_vect = bert_vectorizer(X, verbose=True)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "F4_pN0PICAPu",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "X_square_norm = np.sum(X_vect**2, axis=1).reshape((-1,1))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QiQOuRwAS7ry",
        "colab_type": "text"
      },
      "source": [
        "On CPU, using precomputed norm yields a ~4X improvement over the non-optimized implementation."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "F4Gd4TZEVOJ9",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "test_id = 10060 # The Matrix\n",
        "\n",
        "retriever = L2Retriever(768, use_norm=True, top_k=10, use_gpu=False)\n",
        "%timeit retriever.predict(X_vect, X_vect[test_id], X_square_norm)[0][1:]"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "N9AuVbwT-LKb",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "retriever = L2Retriever(768, use_norm=False, top_k=10, use_gpu=False)\n",
        "%timeit retriever.predict(X_vect, X_vect[test_id])[0][1:]"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LCJOkYZB-LNq",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "%timeit euclidean_distances(X_vect, X_vect[test_id:test_id+1])"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "e0GIAiXoR_uX",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BeEtcocNR_08",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def buildMovieRecommender(movie_names, vectorized_plots, top_k=10):\n",
        "  retriever = L2Retriever(vectorized_plots.shape[1], use_norm=True, top_k=top_k, use_gpu=False)\n",
        "  vectorized_norm = np.sum(vectorized_plots**2, axis=1).reshape((-1,1))\n",
        "  \n",
        "  def recommend(query):\n",
        "    try:\n",
        "      idx = retriever.predict(vectorized_plots, \n",
        "                              vectorized_plots[movie_names.index(query)], \n",
        "                              vectorized_norm)[0][1:]\n",
        "      for i in idx:\n",
        "        print(names[i])\n",
        "    except ValueError:\n",
        "      print(\"{} not found in movie db. Suggestions:\")\n",
        "      for i, name in enumerate(movie_names):\n",
        "        if query.lower() in name.lower():\n",
        "          print(i, name)\n",
        "          \n",
        "  return recommend"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_1knPUoX-LTb",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "recommend = buildMovieRecommender(names, X_vect)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UbT-2kOlG-y3",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "recommend(\"The Matrix\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "R9LA16PNMeNK",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "recommend(\"Terminator 2: Judgment Day\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "stf5OaC4hOmh",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "recommend(\"Star Wars Episode V: The Empire Strikes Back\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jPB78fptMeQf",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "recommend(\"Seven Samurai\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "q3AFgYaLYIf0",
        "colab_type": "text"
      },
      "source": [
        "Even without supervision, the model performs adequately on classification and retrieval tasks. While the model performance can be improved with supervised data, the described approach to text feature extraction provides a solid baseline for downstream NLP solutions. \n",
        "\n",
        "This conludes the guide to building a search engine with BERT and Tensorflow.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KbyKQZF2dMBC",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}